#!/usr/bin/python
# -*- coding: UTF-8 -*-

# Probablistic Road Map (PRM)  随机路标图
import random
import math
import numpy as np
import scipy.spatial
import matplotlib.pyplot as plt
# parameter
N_SAMPLE = 500              # number of sample_points
N_KNN = 10                  # number of edge from one sampled point
MAX_EDGE_LEN = 30.0         # [m] Maximum edge length

show_animation = True

# Node class for Dijkstra search
class Node:
    def __init__(self, x, y, cost, pind):
        self.x = x
        self.y = y
        self.cost = cost
        self.pind = pind    # 父节点
    def __str__(self):
        return str(self.x) + "," + str(self.y) + "," + str(self.cost) + "," + str(self.pind)

# Nearest neighbor search class with KDTree
class KDTree:
    def __init__(self, data):
        # store kd-tree
        self.tree = scipy.spatial.cKDTree(data)
    # nearest neighbor search
    def search(self, inp, k=1):
        if len(inp.shape) >= 2:  # multi input
            index = []
            dist = []
            for i in inp.T:
                idist, iindex = self.tree.query(i, k=k)
                index.append(iindex)
                dist.append(idist)
            return index, dist
        else:
            dist, index = self.tree.query(inp, k=k)
            return index, dist
    # find points with in a distance r
    def search_in_distance(self, inp, r):
        index = self.tree.query_ball_point(inp, r)
        return index

# 碰撞检查
def is_collision(sx, sy, gx, gy, rr, okdtree):
    x = sx
    y = sy
    dx = gx - sx
    dy = gy - sy
    yaw = math.atan2(gy - sy, gx - sx)
    d = math.sqrt(dx**2 + dy**2)
    if d >= MAX_EDGE_LEN:
        return True
    D = rr
    nstep = round(d / D)
    for i in range(nstep):
        idxs, dist = okdtree.search(np.matrix([x, y]).T)
        if dist[0] <= rr:
            return True  # collision
        x += D * math.cos(yaw)
        y += D * math.sin(yaw)
    # goal point check
    idxs, dist = okdtree.search(np.matrix([gx, gy]).T)
    if dist[0] <= rr:
        return True  # collision
    return False  # OK

# 路标图
def plot_road_map(road_map, sample_x, sample_y):
    for i in range(len(road_map)):
        for ii in range(len(road_map[i])):
            ind = road_map[i][ii]
            plt.plot([sample_x[i], sample_x[ind]],[sample_y[i], sample_y[ind]], "-k")
    plt.show()

"""
随机撒点
    1. 必须是自由空间的随机点 
    2. 每个点都要确保与障碍物无碰撞
"""
def sample_points(sx, sy, gx, gy, rr, ox, oy, obstacle_kdtree):
    maxx = max(ox)
    maxy = max(oy)
    minx = min(ox)
    miny = min(oy)
    sample_x, sample_y = [], []
    while len(sample_x) <= N_SAMPLE:
        tx = (random.random() - minx) * (maxx - minx)
        ty = (random.random() - miny) * (maxy - miny)
        index, dist = obstacle_kdtree.search(np.matrix([tx, ty]).T) # search in the kdTree
        if dist[0] >= rr:
            sample_x.append(tx)
            sample_y.append(ty)
    # append start and goal point
    sample_x.append(sx)
    sample_y.append(sy)
    sample_x.append(gx)
    sample_y.append(gy)
    return sample_x, sample_y


# 构建一个路径网络图
# Roadmap is a graph G(V, E)  （无向网络图G，其中V代表随机点集，E代表所有可能的两点之间的路径集）
def generate_roadmap(sample_x, sample_y, rr, obstacle_kdtree):
    """
    Road map generation
    sample_x: [m] x positions of sampled points
    sample_y: [m] y positions of sampled points
    rr: Robot Radius[m]
    obstacle_kdtree: KDTree object of obstacles
    """
    road_map = []
    nsample = len(sample_x)
    # 采样点的kdtree
    sample_points_kdtree = KDTree(np.vstack((sample_x, sample_y)).T)
    for (i, ix, iy) in zip(range(nsample), sample_x, sample_y): # for every (sample_x, sample_y)
        # 寻找临近点的索引  // 临近的k个点
        index, dists = sample_points_kdtree.search(np.matrix([ix, iy]).T, k=nsample)
        inds = index[0][0]
        edge_id = []
        for ii in range(1, len(inds)):
            nx = sample_x[inds[ii]]
            ny = sample_y[inds[ii]]
            if not is_collision(ix, iy, nx, ny, rr, obstacle_kdtree):
                edge_id.append(inds[ii])
            if len(edge_id) >= N_KNN: # number of edge from one sampled point
                break
        road_map.append(edge_id)
    return road_map


def Dijkstra_planning(sx, sy, gx, gy, ox, oy, rr, road_map, sample_x, sample_y):
    """
    gx: goal x position [m]
    gx: goal x position [m]
    ox: x position list of Obstacles [m]
    oy: y position list of Obstacles [m]
    reso: grid resolution [m]
    rr: robot radius[m]
    """
    nstart = Node(sx, sy, 0.0, -1)
    ngoal = Node(gx, gy, 0.0, -1)
    # the open_set and close_set 
    openset = dict()
    closedset = dict()
    openset[len(road_map) - 2] = nstart
    while True:
        # 如果open_set为空
        if len(openset) == 0:
            print("Cannot find path")
            break
        # 取open_set中cost最小的node
        c_id = min(openset, key=lambda o: openset[o].cost)
        current = openset[c_id]
        # show graph
        if show_animation and len(closedset.keys()) % 2 == 0:
            plt.plot(current.x, current.y, "xg")
            plt.pause(0.001)
        if c_id == (len(road_map) - 1):
            print("goal is found!")
            ngoal.pind = current.pind
            ngoal.cost = current.cost
            break
        # 从open_set中删除
        del openset[c_id]
        # 加入closed_set
        closedset[c_id] = current
        # expand search grid based on motion model
        for i in range(len(road_map[c_id])):
            # 加入和该node联通的node
            n_id = road_map[c_id][i]
            dx = sample_x[n_id] - current.x
            dy = sample_y[n_id] - current.y
            d = math.sqrt(dx**2 + dy**2)
            node = Node(sample_x[n_id], sample_y[n_id],current.cost + d, c_id)
            if n_id in closedset:
                continue
            # 是否在open_set
            if n_id in openset:
                if openset[n_id].cost > node.cost:
                    openset[n_id].cost = node.cost
                    openset[n_id].pind = c_id
            else:
                openset[n_id] = node
    # generate final course
    rx, ry = [ngoal.x], [ngoal.y]
    pind = ngoal.pind
    while pind != -1:
        n = closedset[pind]
        rx.append(n.x)
        ry.append(n.y)
        pind = n.pind
    return rx, ry


def PRM_planning(sx, sy, gx, gy, ox, oy, rr):
    # 障碍物的kdtree
    obstacle_kdtree = KDTree(np.vstack((ox, oy)).T)
    # 随机撒点
    sample_x, sample_y = sample_points(sx, sy, gx, gy, rr, ox, oy, obstacle_kdtree) 
    if show_animation:
        plt.plot(sample_x, sample_y, ".b")
    # 产生路标图
    road_map = generate_roadmap(sample_x, sample_y, rr, obstacle_kdtree)
    # 搜索
    rx, ry = Dijkstra_planning( sx, sy, gx, gy, ox, oy, rr, road_map, sample_x, sample_y)
    return rx, ry


def main():
    print(__file__ + " start!!")
    # start and goal position
    sx = 10.0  # [m]
    sy = 10.0  # [m]
    gx = 50.0  # [m]
    gy = 50.0  # [m]
    robot_size = 5.0  # [m]
    # obstacles 
    ox = []
    oy = []
    for i in range(60):
        ox.append(i)
        oy.append(0.0)
    for i in range(60):
        ox.append(60.0)
        oy.append(i)
    for i in range(61):
        ox.append(i)
        oy.append(60.0)
    for i in range(61):
        ox.append(0.0)
        oy.append(i)
    for i in range(40):
        ox.append(20.0)
        oy.append(i)
    for i in range(40):
        ox.append(40.0)
        oy.append(60.0 - i)
    if show_animation:
        plt.plot(ox, oy, ".k")
        plt.plot(sx, sy, "^r")
        plt.plot(gx, gy, "^r")
        plt.grid(True)
        plt.axis("equal")
    # the main process
    rx, ry = PRM_planning(sx, sy, gx, gy, ox, oy, robot_size)
    assert len(rx) != 0, 'Cannot found path'
    if show_animation:
        plt.plot(rx, ry, "-r")
        plt.show()

if __name__ == '__main__':
    main()
